/**
 * @file main.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG//内置宏ASCENDC_CPU_DEBUG是区分运行CPU模式或NPU模式逻辑的标志
#include "acl/acl.h"
extern void sync_bn_training_update_do(uint32_t blockDim, void *stream, 
                                      uint8_t *local_mean, uint8_t *local_var,
                                      uint8_t *global_mean, uint8_t *global_var,
                                      uint8_t *updated_mean, uint8_t *updated_var,
                                      float momentum, float epsilon);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void sync_bn_training_update(GM_ADDR local_mean, GM_ADDR local_var, 
                                                             GM_ADDR global_mean, GM_ADDR global_var,
                                                             GM_ADDR updated_mean, GM_ADDR updated_var,
                                                             float momentum, float epsilon);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;
    size_t tensorByteSize = 2048 * sizeof(float); // 每个张量的大小
    float momentum = 0.9f;
    float epsilon = 1e-5f;

// 用于CPU模式调试的调用程序
#ifdef ASCENDC_CPU_DEBUG
    uint8_t *local_mean = (uint8_t *)AscendC::GmAlloc(tensorByteSize);
    uint8_t *local_var = (uint8_t *)AscendC::GmAlloc(tensorByteSize);
    uint8_t *global_mean = (uint8_t *)AscendC::GmAlloc(tensorByteSize);
    uint8_t *global_var = (uint8_t *)AscendC::GmAlloc(tensorByteSize);
    uint8_t *updated_mean = (uint8_t *)AscendC::GmAlloc(tensorByteSize);
    uint8_t *updated_var = (uint8_t *)AscendC::GmAlloc(tensorByteSize);

    ReadFile("./input/input_local_mean.bin", tensorByteSize, local_mean, tensorByteSize);
    ReadFile("./input/input_local_var.bin", tensorByteSize, local_var, tensorByteSize);
    ReadFile("./input/input_global_mean.bin", tensorByteSize, global_mean, tensorByteSize);
    ReadFile("./input/input_global_var.bin", tensorByteSize, global_var, tensorByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(sync_bn_training_update, blockDim, local_mean, local_var, global_mean, global_var, 
                updated_mean, updated_var, momentum, epsilon); // use this macro for cpu debug

    WriteFile("./output/output_updated_mean.bin", updated_mean, tensorByteSize);
    WriteFile("./output/output_updated_var.bin", updated_var, tensorByteSize);

    AscendC::GmFree((void *)local_mean);
    AscendC::GmFree((void *)local_var);
    AscendC::GmFree((void *)global_mean);
    AscendC::GmFree((void *)global_var);
    AscendC::GmFree((void *)updated_mean);
    AscendC::GmFree((void *)updated_var);
#else// NPU模式运行算子的调用程序
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *local_mean_host, *local_var_host, *global_mean_host, *global_var_host, *updated_mean_host, *updated_var_host;
    uint8_t *local_mean_device, *local_var_device, *global_mean_device, *global_var_device, *updated_mean_device, *updated_var_device;

    // 分配主机内存
    CHECK_ACL(aclrtMallocHost((void **)(&local_mean_host), tensorByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&local_var_host), tensorByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&global_mean_host), tensorByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&global_var_host), tensorByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&updated_mean_host), tensorByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&updated_var_host), tensorByteSize));
    
    // 分配设备内存
    CHECK_ACL(aclrtMalloc((void **)&local_mean_device, tensorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&local_var_device, tensorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&global_mean_device, tensorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&global_var_device, tensorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&updated_mean_device, tensorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&updated_var_device, tensorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    // 读取输入文件
    ReadFile("./input/input_local_mean.bin", tensorByteSize, local_mean_host, tensorByteSize);
    ReadFile("./input/input_local_var.bin", tensorByteSize, local_var_host, tensorByteSize);
    ReadFile("./input/input_global_mean.bin", tensorByteSize, global_mean_host, tensorByteSize);
    ReadFile("./input/input_global_var.bin", tensorByteSize, global_var_host, tensorByteSize);

    // 拷贝数据到设备
    CHECK_ACL(aclrtMemcpy(local_mean_device, tensorByteSize, local_mean_host, tensorByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(local_var_device, tensorByteSize, local_var_host, tensorByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(global_mean_device, tensorByteSize, global_mean_host, tensorByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(global_var_device, tensorByteSize, global_var_host, tensorByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    // 调用算子
    sync_bn_training_update_do(blockDim, stream, 
                              local_mean_device, local_var_device, 
                              global_mean_device, global_var_device, 
                              updated_mean_device, updated_var_device, 
                              momentum, epsilon);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    // 拷贝结果回主机
    CHECK_ACL(aclrtMemcpy(updated_mean_host, tensorByteSize, updated_mean_device, tensorByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    CHECK_ACL(aclrtMemcpy(updated_var_host, tensorByteSize, updated_var_device, tensorByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    
    // 写入输出文件
    WriteFile("./output/output_updated_mean.bin", updated_mean_host, tensorByteSize);
    WriteFile("./output/output_updated_var.bin", updated_var_host, tensorByteSize);

    // 释放资源
    CHECK_ACL(aclrtFree(local_mean_device));
    CHECK_ACL(aclrtFree(local_var_device));
    CHECK_ACL(aclrtFree(global_mean_device));
    CHECK_ACL(aclrtFree(global_var_device));
    CHECK_ACL(aclrtFree(updated_mean_device));
    CHECK_ACL(aclrtFree(updated_var_device));
    
    CHECK_ACL(aclrtFreeHost(local_mean_host));
    CHECK_ACL(aclrtFreeHost(local_var_host));
    CHECK_ACL(aclrtFreeHost(global_mean_host));
    CHECK_ACL(aclrtFreeHost(global_var_host));
    CHECK_ACL(aclrtFreeHost(updated_mean_host));
    CHECK_ACL(aclrtFreeHost(updated_var_host));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}